package leetcode.weekly.week346;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;

//Solution4Test
public class Solution4 {

	class Node {
		int from;
		int dist;
		int index;
		List<Integer> edges;

		public Node(int index, int from) {
			this.index = index;
			this.from = from;
			dist = 0;
			edges = new ArrayList<>();
		}

	}

	public int[][] modifiedGraphEdges(int n, int[][] edges, int source, int destination, int target) {
		int ml = edges.length, nl = edges[0].length;
		List<List<Integer>> nexts = new ArrayList<>();
		for (int i = 0; i < n; i++) {
			nexts.add(new ArrayList<>());
		}
		int[][] ws = new int[n][n];
		for (int i = 0; i < edges.length; i++) {
			int u = edges[i][0], v = edges[i][1];
			nexts.get(u).add(v);
			nexts.get(v).add(u);
			ws[u][v] = edges[i][2];
			ws[v][u] = edges[i][2];
		}
		PriorityQueue<Node> queue = new PriorityQueue<>((a, b) -> a.dist - b.dist);
		int[][] been = new int[n][2];
		for(int i =0;i<n;i++) {
			Arrays.fill(been[i], Integer.MAX_VALUE);
		}
		Node[] nodes = new Node[n];
		for (int i = 0; i < n; i++) {
			nodes[i] = new Node(i, -1);
		}
		queue.add(nodes[source]);
		been[source][0] = 0;
		int[][] ans = null;
		while (!queue.isEmpty()) {
			Node cur = queue.poll();
			if (cur.dist > target) {
				continue;
			}
			if (cur.index == destination) {
				if (cur.dist == target) {
					for (int[] e : edges) {
						if (e[2] == -1) {
							e[2] = 1;
						}
					}
					return edges;
				}
				if (cur.dist < target && cur.edges.size() == 0) {
					return new int[][] {};
				}
				if (cur.dist < target && cur.edges.size() > 0 && ans == null) {
					ans = new int[ml][];
					boolean add = false;
					int i = 0;
					for (int[] e : edges) {
						ans[i] = e;
						if (ans[i][2] == -1) {
							ans[i][2] = 1;
						}
						if (!add && cur.edges.contains(f(e[0], e[1]))) {
							add = true;
							ans[i][2] += target - cur.dist;
						}
						i++;
					}
				}
				continue;
			}
			boolean neg = cur.edges.size() > 0;
			List<Integer> nxts = nexts.get(cur.index);
			for (int nxt : nxts) {
				if (nxt == cur.from) {
					continue;
				}
				int w = ws[cur.index][nxt];
				int y = neg || w == -1 ? 1 : 0;
				if (cur.dist + (w == -1 ? 1 : w) <been[nxt][y]) {
					been[nxt][y] = cur.dist + (w == -1 ? 1 : w);
					Node nxtNode = new Node(nxt, cur.index);
					nxtNode.dist = cur.dist + (w == -1 ? 1 : w);
					nxtNode.edges.addAll(cur.edges);
					if (w == -1) {
						nxtNode.edges.add(f(cur.index, nxt));
					}
					queue.add(nxtNode);
				}
			}

		}
		return ans;
	}

	private int f(int a, int b) {
		if (a < b)
			return a * 200 + b;
		return b * 200 + a;
	}
}
